from threading import Thread
from time import sleep
#import hashlib
import zipfile
import socket
import time
import sys
import re
import os

global s
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

# if_dir2zip_done用于等待dir2zip多线程函数的完成
# if_remove_dir2zip用于声明是否是目录+是否在完成传输后删除临时文件zip
def cast_file(file):
    if_remove_dir2zip = False
    if os.path.exists(file)==False:
        print('!! Cannot Reach '+file)
        sys.exit(0)
    elif os.path.isdir(file):
        if_remove_dir2zip = True
        dir2zip.if_dir2zip_done = False
        try:
            os.remove(file+'.zip')
        except:pass
        dir2zip.file=file
        t=Thread(target=dir2zip)
        t.start()
    else:
        dir2zip.if_dir2zip_done = True

    print("... Waiting for Connection", end='\r')
    #print(hashlib.md5(open(file, "rb").read()).hexdigest())

    s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR, 1)
    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    s.bind(('', 1079))
    s.listen(10)
    c, addr = s.accept()
    print('{} connected.'.format(addr))

    while dir2zip.if_dir2zip_done==False:
        # 如果不加sleep，那么while-pass导致的io阻塞会很严重
        sleep(0.5)
        pass

    # 对接收者告知是否是目录
    if if_remove_dir2zip == True:
        if_remove_dir2zip_signal = b'0x01'
        file = "%s.zip" %file
    else:
        if_remove_dir2zip_signal = b'0x02'
    print('Cast ... ( ', file, ' )')
    c.send(os.path.basename(file).encode('utf-8')+b'0x00'+if_remove_dir2zip_signal+b'0x00'+str(os.path.getsize(file)).encode('utf-8')+b'0x00')
    with open(file, "rb") as open_file:
        for piece in read_inPieces(open_file):
            c.sendall(piece)
    s.close()

    print("Cast : Done")

    if if_remove_dir2zip == True:
        os.remove(file)

def connect_to(env_ip, env_input):
    if env_ip == None:
        target_ip = env_input
    else:
        env_ip_a = env_ip.split('.')
        env_ip_a[3] = ''
        env_ip_a = '.'.join(env_ip_a)
        
        target_ip = env_input.split('.')
        if len(target_ip) == 1:
            target_ip = "%s%s" %(env_ip_a,target_ip[0])
        elif len(target_ip) == 4:
            target_ip = '.'.join(target_ip)
    s.connect((target_ip, 1079))

def dir2zip():
    import shutil
    # dir2zip.file = ''
    # dir2zip.if_dir2zip_done = True
    shutil.make_archive(dir2zip.file, 'zip', dir2zip.file)
    dir2zip.file = dir2zip.file+'.zip'
    dir2zip.if_dir2zip_done = True
# 声明压缩目录的多线程函数

# 取得本地IP地址
def get_env_ip():
    try:
        env_ip_fetch = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        env_ip_fetch.settimeout(2)
        env_ip_fetch.connect(("223.5.5.5",53))
        env_ip = env_ip_fetch.getsockname()[0]
        env_ip_fetch.close()
        return env_ip
    except:
        pass        

def get_file():
    get_file.if_stop_get=False
    get_file.list_file_data = []
    get_file.file_size_total = 0
    get_file.file_get_name = ''
    get_file.file_get = ''
    get_file.get_file_size_print = ''

    t=Thread(target=write_file)
    t2=Thread(target=get_file_size)

    get = s.recv(131072)
    if_create_file=False
    t.start()
    while get:
        if if_create_file==False:
            get=get.split(b'0x00',3)
            try:
                os.remove(get[0])
            except:
                pass
            get_file.file_get_name = get[0].decode(encoding="utf-8")
            get_file.file_get = open(get_file.file_get_name, "ab+")
            if_create_file=True
            if get[1] == b'0x01':
                if_unzip = True
                print("Get : <dir> %s" %get_file.file_get_name.split('.zip',-1)[0])
            elif get[1] == b'0x02':
                if_unzip = False
                print("Get : %s" %get_file.file_get_name)
            get_file.file_size_total = int(get[2].decode(encoding="utf-8"))
            t2.start()
            get_file.list_file_data.append(get[3])
        #print(get)
        get = s.recv(131072)
        get_file.list_file_data.append(get)

    s.close()
    while get_file.list_file_data!=[]:
        sleep(0.5)
    # write_file多线程函数将在声明if_stop_get=True + 数据列表为空后结束
    # 用空格覆盖get_file_size多线程函数的print内容
    print(''.ljust(len(get_file.get_file_size_print)),end='\r')
    get_file.if_stop_get=True
    get_file.file_get.close()

    #print(hashlib.md5(open(f_name[0], "rb").read()).hexdigest())
    print("Get : Done")

    if if_unzip == True:
        dir_name = os.getcwd() + '//' + get_file.file_get_name.split('.zip',-1)[0]
        print_line = "Extract : <dir> %s" %get_file.file_get_name.split('.zip',-1)[0]
        print(print_line, end='\r')
        try:
            os.mkdir(dir_name)
        except:pass
        with zipfile.ZipFile(get_file.file_get_name,'r') as file_as_zip:
            file_as_zip.extractall(path=dir_name)
            file_as_zip.close()
        try:
            os.remove(get_file.file_get_name)
        except:
            pass
        print(''.ljust(len(print_line)),end='\r')

def get_file_size():
    file_size_total=get_file.file_size_total
    file_get_name=get_file.file_get_name
    get_file.if_stop_get

    while get_file.if_stop_get==False:
        get_file.get_file_size_print = '{}{:-2}{}{}{}'.format('Get : ',int(os.path.getsize(file_get_name)/1000000),' / ',file_size_total/1000000,' MB')
        print(get_file.get_file_size_print,end='\r')
        sleep(1)

# 保证读取文件时较少的内存占用
def read_inPieces(read_inPieces_file, read_inPieces_size=131072):
    while True:
        read_inPieces_piece = read_inPieces_file.read(read_inPieces_size)
        if not read_inPieces_piece:
            break
        yield read_inPieces_piece

# 依赖于数据的列表+写入的文件
def write_file():

    while True:
        if get_file.if_stop_get==True and get_file.list_file_data==[]:
            break
        # 在数据列表为空+不打扰在数据列表中加入数据+传输完成时：写入数据到文件
        if get_file.list_file_data!=[]:
            if len(get_file.list_file_data)>= 2 or get_file.if_stop_get==False:
                try:
                    get_file.file_get.write(get_file.list_file_data[0])
                    get_file.list_file_data.pop(0)
                except:
                    pass

def main(args):
    env_ip = get_env_ip()
    print('Local IP is ', env_ip ,end='\n\n')
    if len(args)==2:
        env_input = args[1]
    elif len(args)==1:
        env_input = input('Cast file path, or Get from IP : ')
    else:
        env_input = args[1:]
        f_name = 'transmission_%s.zip' %(time.time())
        f=zipfile.ZipFile(f_name, 'w')
        for env_input_i in env_input:
            f.write(env_input_i, os.path.basename(env_input_i))
        env_input = os.path.abspath(f_name)
        f.close()
    env_input = env_input.strip('\'\"\\').rstrip("/")
    if not os.path.exists(env_input):
        try:
            connect_to(env_ip,env_input)
        except:
            print("Cannot reach IP ",env_input)
        get_file()
    else:
        cast_file(env_input)

if __name__ == "__main__":
    main(sys.argv)
